#!/usr/bin/env Rscript

rm(list=ls())
library(FactorHet)
library(dplyr)

args <- commandArgs(trailingOnly = TRUE)
slurm_id <- as.numeric(args[1])

repdata_nonhisp <- readRDS('packaged_data.RDS')

set.seed(slurm_id)

id <- sample(unique(repdata_nonhisp$CaseID), 
	length(unique(repdata_nonhisp$CaseID)) * 0.5)

moderator_fmla <- ~ scale_hisp_prej_flip + party_ID + 
     ppEducat + census_div + ppEthm

fmla <- Chosen_Immigrant ~ Ed + Gender + Country + Reason + Job + Exp + 
	  Plans + Trips + Lang + (Job + Ed + Reason) * Country +Job*Ed

current_dir <- getwd()

setwd(tempdir())

output <- list()

for (FACTORHET_K in c(2,3)){

	time_main <- proc.time()
	est_mbo <- FactorHet_mbo(formula = fmla, K = FACTORHET_K,
	   design = repdata_nonhisp %>% filter(CaseID %in% id), 
	   moderator =  moderator_fmla,
	   group = ~ CaseID, task = ~ contest_no, choice_order= ~ choice_id)
	time_main <- proc.time() - time_main


	time_refit <- proc.time()
	refit_mbo <- FactorHet_refit(
		object = est_mbo,
		newdata = repdata_nonhisp %>% filter(!(CaseID %in% id)))
	time_refit <- proc.time() - time_refit

	AME_refit <- AME(refit_mbo, plot = FALSE, design = repdata_nonhisp)
	moderator_split <- margeff_moderators(est_mbo)$data
	output[[FACTORHET_K]] <- list(
		AME_refit = AME_refit,
		moderator = moderator_split,
		information_criterion = est_mbo$information_criterion,
		K = FACTORHET_K
	)

}

setwd(current_dir)
saveRDS(output, 'out.RDS')